#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
人体检测模块 - 使用 YOLOv8 检测照片中是否有人
"""

import os
from pathlib import Path
from typing import List, Tuple, Optional
from PIL import Image
import threading


class PersonDetector:
    """使用 YOLOv8 检测照片中的人体"""

    def __init__(self, confidence_threshold: float = 0.25):
        """
        初始化人体检测器

        Args:
            confidence_threshold: 置信度阈值，默认0.25（越高越严格）
        """
        self.confidence_threshold = confidence_threshold
        self.model = None
        self._load_model()

    def _load_model(self):
        """加载 YOLOv8 模型"""
        try:
            from ultralytics import YOLO
            # 使用预训练的 YOLOv8n 模型（最快的版本）
            # 首次运行会自动下载模型文件
            self.model = YOLO('yolov8n.pt')
            print("YOLOv8 人体检测模型加载成功")
        except Exception as e:
            print(f"加载 YOLOv8 模型失败: {e}")
            self.model = None

    def detect_person_in_image(self, image_path: str) -> Tuple[bool, int, float]:
        """
        检测单张图片中是否有人

        Args:
            image_path: 图片路径

        Returns:
            (是否有人, 检测到的人数, 最高置信度)
        """
        if self.model is None:
            raise RuntimeError("YOLOv8 模型未成功加载")

        try:
            # 运行检测（verbose=False 静默模式）
            results = self.model(image_path, verbose=False)

            # 获取第一个结果（单张图片）
            result = results[0]

            # 获取所有检测框
            boxes = result.boxes

            # 过滤出"人"类别（class 0 是 person）
            person_count = 0
            max_confidence = 0.0

            for box in boxes:
                # 获取类别ID和置信度
                cls = int(box.cls[0])
                conf = float(box.conf[0])

                # 类别0是人，且置信度大于阈值
                if cls == 0 and conf >= self.confidence_threshold:
                    person_count += 1
                    max_confidence = max(max_confidence, conf)

            has_person = person_count > 0
            return has_person, person_count, max_confidence

        except Exception as e:
            print(f"检测图片 {os.path.basename(image_path)} 时出错: {e}")
            return False, 0, 0.0

    def classify_photos(
        self,
        photo_paths: List[str],
        progress_callback: Optional[callable] = None
    ) -> Tuple[List[str], List[str]]:
        """
        批量分类照片：有人 vs 无人

        Args:
            photo_paths: 照片路径列表
            progress_callback: 进度回调函数 callback(current, total, filename, has_person)

        Returns:
            (有人的照片列表, 无人的照片列表)
        """
        with_person = []
        without_person = []

        total = len(photo_paths)

        for idx, photo_path in enumerate(photo_paths, 1):
            try:
                has_person, person_count, confidence = self.detect_person_in_image(photo_path)

                if has_person:
                    with_person.append(photo_path)
                else:
                    without_person.append(photo_path)

                # 调用进度回调
                if progress_callback:
                    progress_callback(idx, total, os.path.basename(photo_path), has_person, person_count, confidence)

            except Exception as e:
                print(f"处理照片 {photo_path} 时出错: {e}")
                # 出错的照片归类到"无人"
                without_person.append(photo_path)

                if progress_callback:
                    progress_callback(idx, total, os.path.basename(photo_path), False, 0, 0.0)

        return with_person, without_person

    def get_model_info(self) -> str:
        """获取模型信息"""
        if self.model is None:
            return "模型未加载"
        return f"YOLOv8n (置信度阈值: {self.confidence_threshold})"


# 测试代码
if __name__ == "__main__":
    detector = PersonDetector(confidence_threshold=0.3)
    print(f"模型信息: {detector.get_model_info()}")

    # 测试单张图片
    test_image = "test.jpg"
    if os.path.exists(test_image):
        has_person, count, conf = detector.detect_person_in_image(test_image)
        print(f"图片: {test_image}")
        print(f"是否有人: {has_person}")
        print(f"人数: {count}")
        print(f"最高置信度: {conf:.2f}")
